"""
Kardar-Parisi-Zhang (KPZ) Interface PDE implementation

The KPZ equation describes the evolution of a growing interface:
    ∂h/∂t = ν∇²h + (λ/2)|∇h|² + η

This is a second-order nonlinear PDE with optional noise.
"""

using ..Fields
using ..Grids
using Random

"""
    KPZInterfacePDE <: AbstractPDE

Kardar-Parisi-Zhang equation for interface growth dynamics.

The mathematical definition is:
    ∂h/∂t = ν∇²h + (λ/2)|∇h|² + η(r,t)

where h is the interface height, ν is the diffusion coefficient, 
λ is the nonlinear coupling strength, and η is Gaussian white noise.

# Parameters
- `nu`: Diffusion coefficient (default: 0.5)
- `lambda`: Nonlinear coupling strength (default: 1.0)  
- `noise`: Noise variance (default: 0.0)
- `bc`: Boundary conditions (default: "auto_periodic_neumann")

# Properties
- Second-order nonlinear PDE
- Can exhibit roughening and scaling behavior
- With noise: stochastic PDE
- Without noise: deterministic interface equation
"""
struct KPZInterfacePDE <: AbstractPDE
    nu::Float64
    lambda::Float64
    noise::Float64
    bc::String
    rng::Random.AbstractRNG
    
    function KPZInterfacePDE(; nu=0.5, lambda=1.0, noise=0.0, 
                           bc="auto_periodic_neumann", rng=Random.default_rng())
        new(nu, lambda, noise, bc, rng)
    end
end

"""
    evolution_rate(pde::KPZInterfacePDE, state::ScalarField, t::Float64)

Compute the evolution rate for the KPZ equation.

The equation is: ∂h/∂t = ν∇²h + (λ/2)|∇h|² + η
"""
function evolution_rate(pde::KPZInterfacePDE, state::ScalarField, t::Float64)
    h = data(state)
    grid_obj = grid(state)
    
    # Linear diffusion term: ν∇²h
    h_laplace = laplacian(state)
    diffusion_term = pde.nu * h_laplace
    
    # Nonlinear gradient term: (λ/2)|∇h|²
    grad_h = gradient(state)
    
    # Compute |∇h|² for each grid point
    grad_magnitude_sq = zeros(size(h))
    if isa(grad_h, VectorField)
        # For VectorField, sum squares of components
        grad_data = data(grad_h)
        if ndims(grad_data) >= 2  # Multi-dimensional gradients
            for i in 1:size(grad_data, 1)  # Loop over spatial dimensions
                grad_magnitude_sq .+= grad_data[i, :].^2
            end
        else
            grad_magnitude_sq = grad_data.^2
        end
    else
        # Handle case where gradient returns array of arrays or similar
        for component in grad_h
            grad_magnitude_sq .+= component.^2
        end
    end
    
    nonlinear_term = ScalarField(pde.lambda * grad_magnitude_sq / 2, grid_obj)
    
    # Combine terms
    result = diffusion_term + nonlinear_term
    
    # Add noise if specified
    if pde.noise > 0
        noise_data = sqrt(pde.noise) * randn(pde.rng, size(h))
        noise_term = ScalarField(noise_data, grid_obj)
        result = result + noise_term
    end
    
    return result
end

"""
    has_noise(pde::KPZInterfacePDE)

Check if the PDE includes stochastic noise terms.
"""
has_noise(pde::KPZInterfacePDE) = pde.noise > 0

"""
    is_nonlinear(pde::KPZInterfacePDE)

Check if the PDE has nonlinear terms.
"""
is_nonlinear(pde::KPZInterfacePDE) = pde.lambda != 0

"""
    get_roughness(pde::KPZInterfacePDE, state::ScalarField)

Compute the interface roughness: W = √⟨(h - ⟨h⟩)²⟩
"""
function get_roughness(pde::KPZInterfacePDE, state::ScalarField)
    h = data(state)
    h_mean = sum(h) / length(h)
    roughness = sqrt(sum((h .- h_mean).^2) / length(h))
    return roughness
end

"""
    get_interface_velocity(pde::KPZInterfacePDE, state::ScalarField)

Estimate the average interface velocity based on the mean slope.
"""
function get_interface_velocity(pde::KPZInterfacePDE, state::ScalarField)
    grad_h = gradient(state)
    
    # Compute mean gradient magnitude
    grad_data = data(grad_h)
    if isa(grad_h, VectorField)
        mean_grad_sq = 0.0
        if ndims(grad_data) >= 2
            for i in 1:size(grad_data, 1)
                mean_grad_sq += sum(grad_data[i, :].^2) / length(grad_data[i, :])
            end
        else
            mean_grad_sq = sum(grad_data.^2) / length(grad_data)
        end
    else
        mean_grad_sq = sum(grad_data.^2) / length(grad_data)
    end
    
    # KPZ velocity contribution
    return pde.lambda * mean_grad_sq / 2
end

"""
    is_stiff(pde::KPZInterfacePDE)

KPZ equation is generally not stiff unless noise is very large.
"""
is_stiff(pde::KPZInterfacePDE) = pde.noise > 10.0

"""
    requires_high_accuracy(pde::KPZInterfacePDE)

KPZ equation with strong nonlinearity may require higher accuracy.
"""
requires_high_accuracy(pde::KPZInterfacePDE) = abs(pde.lambda) > 5.0

"""
    expression(pde::KPZInterfacePDE)

Return a string representation of the KPZ PDE.
"""
function expression(pde::KPZInterfacePDE)
    terms = String[]
    
    # Diffusion term
    if pde.nu != 0
        nu_str = pde.nu == 1.0 ? "" : "$(pde.nu)"
        push!(terms, "$(nu_str)∇²h")
    end
    
    # Nonlinear term
    if pde.lambda != 0
        lambda_str = pde.lambda == 1.0 ? "0.5" : "$(pde.lambda/2)"
        push!(terms, "$(lambda_str)|∇h|²")
    end
    
    # Noise term
    if pde.noise > 0
        push!(terms, "√$(pde.noise)η")
    end
    
    if isempty(terms)
        return "0"
    else
        return join(terms, " + ")
    end
end

# Export the PDE type and analysis functions
export KPZInterfacePDE, has_noise, is_nonlinear, get_roughness, get_interface_velocity